Skip to content

fix(pcie): isolate nested CUDA graph capture channels - #42

Merged
lukealonso merged 1 commit into
local-inference-lab:masterfrom
voipmonitor:fix/pcie-oneshot-nested-capture-channel-20260719
Jul 19, 2026
Merged

fix(pcie): isolate nested CUDA graph capture channels#42
lukealonso merged 1 commit into
local-inference-lab:masterfrom
voipmonitor:fix/pcie-oneshot-nested-capture-channel-20260719

Conversation

@voipmonitor

@voipmonitor voipmonitor commented Jul 19, 2026

Copy link
Copy Markdown
Contributor

Summary

  • bind torch/Inductor nested CUDA capture streams to the B12X channel selected by their enclosing graph capture
  • keep target and MTP/draft graph managers on distinct PCIe oneshot signal and staging buffers
  • apply the same capture ownership model to the PCIe DCP A2A pool
  • preserve legacy fallback behavior for captures entered without a pool capture context
  • add CPU regression coverage for separate target/draft outer and nested stream pairs in both pools

Root cause

A communication channel owns mutable signal counters and double-buffered IPC staging memory. The old nested-capture fallback selected the first existing channel because allocating during CUDA capture is illegal. Independent target and MTP graph managers could therefore capture one shared channel and replay it concurrently, corrupting its state.

The pools now maintain an enclosing capture-channel stack. A torch/Inductor-owned nested capture stream aliases only the channel selected by its own outer graph manager. Target and draft captures consequently remain isolated without adding a runtime copy or synchronization.

Validation

  • tests/distributed/test_pcie_oneshot.py and tests/distributed/test_pcie_dcp_a2a.py: 32 passed
  • Ruff and git diff --check: passed
  • GLM-5.2 TP6/DCP6/MTP3 with vLLM multichannel wiring survived the previously deterministic CC4 failure for both original and online-MXFP8 configurations

Integration

@coderabbitai

coderabbitai Bot commented Jul 19, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

Caution

Review failed

The pull request is closed.

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 3c1071d7-7469-4f02-8bb3-bb76f43904bb

📥 Commits

Reviewing files that changed from the base of the PR and between 21e508e and 2d710b8.

📒 Files selected for processing (4)
  • b12x/distributed/pcie_dcp_a2a.py
  • b12x/distributed/pcie_oneshot.py
  • tests/distributed/test_pcie_dcp_a2a.py
  • tests/distributed/test_pcie_oneshot.py

Walkthrough

PCIeDCPA2APool and PCIeOneshotAllReducePool now preserve channel selection across nested CUDA graph captures, register reused channels under current stream keys, validate capture-stack integrity, and test separate and nested capture scenarios.

Changes

PCIe capture channel reuse

Layer / File(s) Summary
Pool capture channel behavior
b12x/distributed/pcie_dcp_a2a.py, b12x/distributed/pcie_oneshot.py
Both pools track active capture channels, reuse the enclosing channel during capture, register it under the current stream key, and validate stack push/pop consistency.
Capture channel validation
tests/distributed/test_pcie_dcp_a2a.py, tests/distributed/test_pcie_oneshot.py
Tests verify distinct channels for separate captures and reuse of the outer channel during nested captures.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Caller
  participant PCIePool
  participant ChannelFactory
  Caller->>PCIePool: enter capture(stream)
  PCIePool->>ChannelFactory: create or retrieve channel
  PCIePool->>PCIePool: push channel onto capture stack
  Caller->>PCIePool: for_stream()
  PCIePool->>PCIePool: reuse top captured channel
  PCIePool-->>Caller: return captured channel
  Caller->>PCIePool: exit capture
  PCIePool->>PCIePool: pop and validate channel
Loading
🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title clearly and concisely describes the main change: isolating nested CUDA graph capture channels for PCIe oneshot behavior.
Docstring Coverage ✅ Passed No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@voipmonitor
voipmonitor force-pushed the fix/pcie-oneshot-nested-capture-channel-20260719 branch from 21e508e to 2d710b8 Compare July 19, 2026 16:04
@lukealonso
lukealonso merged commit 00695ee into local-inference-lab:master Jul 19, 2026
1 check was pending
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants